package common

import (
	"context"
	"errors"
	"github.com/go-stomp/stomp"
	"github.com/sirupsen/logrus"
)

type ArtemisMQMsgHandler func(message *stomp.Message)

type ArtemisMQ struct {
	addr       string        //服务器地址
	username   string        //账号
	password   string        //密码
	prefix     string        //topic名称前缀
	log        *logrus.Entry //日志
	conn       *stomp.Conn
	clientSubs []*artemisMQSub //客户端
}

type artemisMQSub struct {
	ctx        context.Context
	cancelFunc context.CancelFunc
	topic      string
	sub        *stomp.Subscription
	handler    ArtemisMQMsgHandler
}

// activeMQ artemis
func NewArtemisMQ(addr string, username string, password string, prefix string) (*ArtemisMQ, error) {
	ret := &ArtemisMQ{
		addr:     addr,
		username: username,
		password: password,
		prefix:   prefix,
	}
	err := ret.Init(false)
	return ret, err
}

func (a *ArtemisMQ) Init(reload bool) error {
	if !reload && a.conn != nil {
		return nil
	}
	a.Close()
	conn, err := stomp.Dial("tcp", a.addr, stomp.ConnOpt.Login(a.username, a.password))
	if err != nil {
		return err
	}
	a.conn = conn
	for _, v := range a.clientSubs { //重新监听
		a.addSubscribe(v)
	}
	return nil
}

func (a *ArtemisMQ) Reload() {
	a.Init(true)
}

//创建客户端
func (a *ArtemisMQ) NewSubscribe(topic string, handler ArtemisMQMsgHandler) error {
	for _, v := range a.clientSubs {
		if v.topic == topic {
			return errors.New("该话题已有订阅")
		}
	}
	req := &artemisMQSub{
		topic:   topic,
		handler: handler,
	}
	err := a.addSubscribe(req)
	if err != nil {
		return err
	}
	a.clientSubs = append(a.clientSubs, req)
	return nil
}

func (a *ArtemisMQ) addSubscribe(req *artemisMQSub) error {
	a.Init(false)
	if a.conn == nil {
		return errors.New("nats未连接")
	} else if req.handler == nil {
		return errors.New("处理函数空")
	}
	sub, err := a.conn.Subscribe(a.prefix+req.topic, stomp.AckAuto)
	if err != nil {
		return err
	}
	req.ctx, req.cancelFunc = context.WithCancel(context.Background())
	go func(ctx context.Context, sub *stomp.Subscription, handler ArtemisMQMsgHandler) {
		for {
			select {
			case msg := <-sub.C:
				handler(msg)
			case <-ctx.Done():
				sub.Unsubscribe()
				return
			}
		}
	}(req.ctx, sub, req.handler)
	req.sub = sub
	return nil
}

//推送一个消息
func (a *ArtemisMQ) Put(topic string, msg []byte) error {
	a.Init(false)
	return a.conn.Send(a.prefix+topic, "text/plain", msg)
}

func (a *ArtemisMQ) Close() {
	for _, v := range a.clientSubs {
		if v.cancelFunc != nil {
			v.cancelFunc()
			v.sub = nil
		}
	}
	if a.conn != nil {
		a.conn.Disconnect()
		a.conn = nil
	}
}
